package cn.com.w.ratelimiter.aspect;

import cn.com.w.ratelimiter.annotation.Limiting;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.concurrent.ConcurrentHashMap;


@Component
@Aspect
@Slf4j
public class RateLimitAspect {

    /**
     * key=方法名 value=速度限制器。
     */
    private ConcurrentHashMap<String, RateLimiter> RATE_LIMITER_MAP = new ConcurrentHashMap<>();

    /**
     * 速度限制器
     */
    private RateLimiter rateLimiter;

    @Pointcut("@annotation(cn.com.w.ratelimiter.annotation.Limiting)")
    public void doLimit() {
    }

    @Around("doLimit()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        // 获取拦截的方法签名 eg: String cn.com.w.ratelimiter.controller.HelloController.hello(String)
        Signature signature = point.getSignature();
        // 获取拦截的方法名 eg: String cn.com.w.ratelimiter.controller.HelloController.hello(String)
        MethodSignature msg = (MethodSignature) signature;
        // 返回被织入增加处理目标对象
        Object target = point.getTarget();
        // 获取方法信息 -> 为了获取注解信息
        Method currentMethod = target.getClass().getMethod(msg.getName(), msg.getParameterTypes());
        // 获取注解信息
        Limiting annotation = currentMethod.getAnnotation(Limiting.class);
        // 获取注解每秒加入桶中的令牌数量(QPS)
        double limitNum = annotation.limitNum();
        // 注解所在方法名区分不同的限流策略
        String functionName = msg.getName();

        if (RATE_LIMITER_MAP.containsKey(functionName)) {
            rateLimiter = RATE_LIMITER_MAP.get(functionName);
            log.info("已缓存方法【{}】限流数据。", functionName);
        } else {
            RATE_LIMITER_MAP.put(functionName, RateLimiter.create(limitNum));
            rateLimiter = RATE_LIMITER_MAP.get(functionName);
            log.info("进行缓存方法【{}】限流数据。", functionName);
        }

        if (rateLimiter.tryAcquire()) {
            log.info("方法【{}】未达到限流上限，继续进行业务处理。", functionName);
            return point.proceed();
        } else {
            StringBuffer sb = new StringBuffer("方法【").append(functionName).append("】达到限流上限。");
            throw new RuntimeException(sb.toString());
        }
    }

}
